#!/bin/bash

clear

# 启用BuildKit
export DOCKER_BUILDKIT=1

# 显示菜单
show_menu() {
    GREEN='\033[0;32m'
    BLUE='\033[0;34m'
    YELLOW='\033[1;33m'
    CYAN='\033[0;36m'
    RED='\033[0;31m'
    NC='\033[0m'
    echo -e "${BLUE}===== Python 服务管理脚本 =====${NC}"
    echo -e "🏠${GREEN}本地服务${NC}:"
    echo -e "  1. ${GREEN}在本地环境运行服务${NC}"
    echo -e "🐳${CYAN}Docker服务${NC}:"
    echo -e "  2. ${CYAN}构建并运行Docker容器${NC}"
    echo -e "  3. ${CYAN}仅构建镜像${NC}"
    echo -e "  4. ${CYAN}启动已构建的容器${NC}"
    echo -e "  5. ${CYAN}停止Docker容器${NC}"
    echo -e "  6. ${CYAN}查看Docker容器日志${NC}"
    echo -e "  7. ${CYAN}清理缓存并完全重建${NC}"
    echo -e "🔧${GREEN}Docker维护${NC}:"
    echo -e "  8. ${GREEN}诊断并修复Docker问题${NC}"
    echo -e "🔍${BLUE}代码检查${NC}:"
    echo -e "  9. ${BLUE}安装代码检查工具${NC}"
    echo -e "  10. ${BLUE}运行代码检查${NC}"
    echo -e "🔍${CYAN}环境检查${NC}:"
    echo -e "  11. ${CYAN}检查Python环境${NC}"
    echo -e "  12. ${CYAN}验证 GPU 可用性${NC}"
    echo -e "📊${GREEN}Docker监控${NC}:"
    echo -e "  13. ${GREEN}监控Docker容器${NC}"
    echo -e "  14. ${GREEN}监控Docker构建进程${NC}"
    echo -e "  0. ${RED}退出${NC}"
    echo -e "${BLUE}==============================${NC}"
    echo -n "请选择操作 [0-14]: "
}


# ===========================================================
# 本地环境运行
# ===========================================================
run_local() {
    # 清屏
    clear

    # 直接调用 start_fastapi.sh 脚本，传递所有参数
    bash start_fastapi.sh "$@"
}


# ===========================================================
# Docker操作 - 构建并启动
# ===========================================================
run_docker() {
    echo "正在构建并启动Docker容器..."

    # 使用详细输出模式构建
    echo "构建镜像..."
    if ! sudo docker compose build; then
        echo "❌ 构建失败！不会启动容器。"
        echo "请检查网络连接或尝试选项8修复Docker环境。"
        return 1
    fi

    # 启动容器
    echo "启动容器..."
    if ! sudo docker compose up -d; then
        echo "❌ 容器启动失败！"
        return 1
    fi

    echo "✅ 容器已启动，服务运行在 ws://localhost:10096"
}

# 仅构建Docker镜像
build_only() {
    echo "正在构建镜像..."
    if ! sudo docker compose build; then
        echo "❌ 构建失败！"
        echo "请检查网络连接或尝试选项8修复Docker环境。"
        return 1
    fi
    echo "✅ 构建完成"
}

# 启动已构建的服务
start_service() {
    echo "正在启动服务..."
    if ! sudo docker compose up -d; then
        echo "❌ 容器启动失败！"
        return 1
    fi
    echo "✅ 服务已启动，访问地址: ws://localhost:10096"
}

# 停止Docker容器
stop_docker() {
    echo "正在停止Docker容器..."
    if ! sudo docker compose down; then
        echo "❌ 容器停止失败！"
        return 1
    fi
    echo "✅ 容器已停止"
}

# 查看Docker日志
show_logs() {
    echo "显示容器日志 (Ctrl+C 退出)..."
    sudo docker compose logs -f
}

# 完全重建Docker容器
rebuild_all() {
    echo "正在清理缓存并完全重建..."

    # 停止容器
    echo "停止现有容器..."
    sudo docker compose down

    # 尝试清理构建缓存
    echo "清理构建缓存..."
    if ! sudo docker builder prune -f; then
        echo "⚠️ 警告: 构建缓存清理失败，将尝试重启Docker服务"
        fix_docker
    fi

    # 强制不使用缓存构建
    echo "开始重新构建镜像..."
    if ! sudo docker compose build --no-cache; then
        echo "❌ 构建失败，尝试修复Docker环境后重试..."
        fix_docker
        echo "再次尝试构建..."
        if ! sudo docker compose build --no-cache; then
            echo "❌ 第二次构建也失败了！不会启动容器。"
            echo "请检查网络连接和Docker环境。"
            return 1
        fi
    fi

    # 启动容器
    echo "启动容器..."
    if ! sudo docker compose up -d; then
        echo "❌ 容器启动失败！"
        return 1
    fi

    echo "✅ 重建完成，服务已启动"
}

# 诊断并修复Docker问题
fix_docker() {
    echo "正在诊断Docker环境..."

    # 检查网络连接
    echo "检查网络连接..."
    ping -c 3 registry-1.docker.io || echo "⚠️ 无法连接到Docker Hub，可能存在网络问题"

    # 检查Docker服务状态
    echo "检查Docker服务状态..."
    if command -v systemctl &> /dev/null; then
        sudo systemctl status docker --no-pager || echo "⚠️ Docker服务可能未正常运行"
    else
        # macOS不使用systemctl
        docker info || echo "⚠️ Docker可能未正常运行"
    fi

    # 检查Docker磁盘空间
    echo "检查Docker磁盘空间..."
    if [ -d "/var/lib/docker" ]; then
        sudo df -h /var/lib/docker
    else
        # macOS的Docker不在标准位置
        echo "在macOS上无法直接检查Docker磁盘空间"
    fi

    # 尝试重启Docker服务
    echo "尝试重启Docker服务..."
    if command -v systemctl &> /dev/null; then
        sudo systemctl restart docker
        sleep 3
    else
        # macOS上重启Docker Desktop
        echo "请手动重启Docker Desktop应用"
        read -p "重启Docker Desktop后按回车继续..."
    fi

    # 清理无用的容器和镜像
    echo "清理无用的容器和镜像..."
    sudo docker system prune -f

    # 检查BuildKit状态
    echo "重置BuildKit状态..."
    if [ -f "/var/lib/docker/buildkit/containerdmeta.db" ]; then
        sudo rm -rf /var/lib/docker/buildkit/containerdmeta.db || true
        sudo rm -rf /var/lib/docker/buildkit/metadata.db || true
    else
        echo "在macOS上无法直接访问BuildKit数据库文件"
        echo "请考虑在Docker Desktop设置中重置Docker"
    fi

    echo "Docker环境诊断和修复完成"
}




# ===========================================================
# 代码检查
# ===========================================================
# 安装代码检查工具
install_code_check_tools() {
    echo "正在安装代码检查工具..."

    # 检查虚拟环境是否存在并激活
    if [ ! -d ".venv" ]; then
        echo "正在创建虚拟环境..."
        uv venv
    fi

    echo "正在激活虚拟环境..."
    source .venv/bin/activate

    # 安装代码检查工具
    echo "正在安装代码检查工具..."
    if ! uv pip install flake8 mypy black isort pylint ruff; then
        echo "❌ 代码检查工具安装失败"
        return 1
    fi

    # 创建 .vscode 目录
    mkdir -p .vscode

    # 创建 VS Code/Cursor 配置文件
    echo "正在创建 VS Code/Cursor 配置文件..."
    cat > .vscode/settings.json << 'EOF'
{
    "python.defaultInterpreterPath": "./.venv/bin/python",
    "python.linting.enabled": true,
    "python.linting.flake8Enabled": true,
    "python.linting.pylintEnabled": true,
    "python.linting.mypyEnabled": true,
    "python.linting.lintOnSave": true,
    "python.linting.flake8Args": [
        "--max-line-length=88",
        "--extend-ignore=E203,W503"
    ],
    "python.linting.mypyArgs": [
        "--ignore-missing-imports",
        "--follow-imports=silent"
    ],
    "python.formatting.provider": "black",
    "python.formatting.blackArgs": [
        "--line-length=88"
    ],
    "python.sortImports.args": [
        "--profile=black"
    ],
    "editor.formatOnSave": true,
    "editor.codeActionsOnSave": {
        "source.organizeImports": "explicit"
    },
    "python.analysis.typeCheckingMode": "strict",
    "python.analysis.autoImportCompletions": true,
    "python.analysis.diagnosticMode": "workspace",
    "python.linting.flake8CategorySeverity.E": "Error",
    "python.linting.flake8CategorySeverity.W": "Warning",
    "python.linting.flake8CategorySeverity.F": "Error",
    "editor.rulers": [88],
    "files.trimTrailingWhitespace": true,
    "files.insertFinalNewline": true
}
EOF

    # 创建 flake8 配置文件
    echo "正在创建 flake8 配置文件..."
    cat > .flake8 << 'EOF'
[flake8]
max-line-length = 88
extend-ignore = E203, W503, E501
exclude =
    .git,
    __pycache__,
    .venv,
    .pytest_cache,
    *.egg-info
per-file-ignores =
    __init__.py:F401
EOF

    # 创建 mypy 配置文件
    echo "正在创建 mypy 配置文件..."
    cat > .mypy.ini << 'EOF'
[mypy]
python_version = 3.9
warn_return_any = True
warn_unused_configs = True
disallow_untyped_defs = True
ignore_missing_imports = True
EOF

    echo "✅ 代码检查工具安装和配置完成！"
    echo "请重新启动 Cursor 编辑器以使配置生效。"
    echo "在编辑器中按 Cmd+Shift+P，搜索 'Python: Select Interpreter'，选择 .venv/bin/python"
}

# 运行代码检查
run_code_check() {
    echo "正在运行代码检查..."

    # 检查虚拟环境
    if [ ! -d ".venv" ]; then
        echo "❌ 虚拟环境不存在，请先运行选项9安装代码检查工具"
        return 1
    fi

    source .venv/bin/activate

    # 检查 Python 文件
    echo "正在查找 Python 文件..."
    python_files=$(find . -name "*.py" -not -path "./.venv/*" -not -path "./__pycache__/*")

    if [ -z "$python_files" ]; then
        echo "❌ 没有找到 Python 文件"
        return 1
    fi

    echo "找到的 Python 文件:"
    echo "$python_files"
    echo ""

    # 运行 flake8 检查
    echo "=== Flake8 代码风格检查 ==="
    if command -v flake8 &> /dev/null; then
        flake8 $python_files
        echo "Flake8 检查完成"
    else
        echo "❌ flake8 未安装"
    fi
    echo ""

    # 运行 mypy 类型检查
    echo "=== MyPy 类型检查 ==="
    if command -v mypy &> /dev/null; then
        mypy $python_files
        echo "MyPy 检查完成"
    else
        echo "❌ mypy 未安装"
    fi
    echo ""

    echo "✅ 代码检查完成！"
}


# ===========================================================
# 环境检查
# ===========================================================
check_python_environment() {
    echo "正在检查Python环境..."
    if command -v python &> /dev/null; then
        python --version
    else
        echo "❌ Python 未安装"
    fi
}

# 检查GPU可用性
check_gpu() {
    echo "正在检查GPU可用性..."
    # 检查MPS是否可用
    if command -v python &> /dev/null; then
        python -c "import torch; print(f'PyTorch版本: {torch.__version__}')"
        python -c "import torch; print(f'MPS available: {torch.backends.mps.is_available()}')"
        python -c "import torch; print(f'MPS built: {torch.backends.mps.is_built()}')"
    else
        echo "❌ Python 未安装"
    fi
}

# ===========================================================
# Docker监控功能
# ===========================================================
# 监控Docker容器
monitor_container() {
    echo "请输入容器名称（留空则自动检测）:"
    read container_name

    # 首先检查Docker是否启动
    if ! sudo docker info >/dev/null 2>&1; then
        echo "❌ Docker服务未启动，请先启动Docker Desktop"
        return 1
    fi

    # 如果用户没有输入容器名，尝试自动检测
    if [ -z "$container_name" ]; then
        echo "正在自动检测运行中的容器..."
        container_name=$(sudo docker ps --format "table {{.Names}}" | grep -v NAMES | head -n 1)
        if [ -z "$container_name" ]; then
            echo "❌ 没有找到运行中的容器"
            return 1
        fi
        echo "检测到容器: $container_name"
    fi

    echo "监控容器: $container_name (输入 'r' 刷新，输入 'q' 退出)"
    echo "----------------------------------------"

    # 首次显示信息
    clear
    # 显示容器状态
    echo "容器状态:"
    sudo docker ps -a | grep $container_name || echo "容器未找到或已停止"

    # 显示容器资源使用情况
    echo -e "\n资源使用情况:"
    if sudo docker ps --format "{{.Names}}" | grep -q "^$container_name$"; then
        sudo docker stats --no-stream $container_name 2>/dev/null || echo "无法获取资源使用情况"
    else
        echo "容器未运行"
    fi

    # 显示最近日志
    echo -e "\n最近日志 (最后10行):"
    sudo docker logs --tail 10 $container_name 2>&1 || echo "无法获取日志"

    echo -e "\n=========================================="
    echo "输入 'r' 刷新，输入 'q' 退出:"

    while true; do
        read -n 1 user_input
        echo  # 换行

        case $user_input in
            r|R)
                clear
                # 显示容器状态
                echo "容器状态:"
                sudo docker ps -a | grep $container_name || echo "容器未找到或已停止"

                # 显示容器资源使用情况
                echo -e "\n资源使用情况:"
                if sudo docker ps --format "{{.Names}}" | grep -q "^$container_name$"; then
                    sudo docker stats --no-stream $container_name 2>/dev/null || echo "无法获取资源使用情况"
                else
                    echo "容器未运行"
                fi

                # 显示最近日志
                echo -e "\n最近日志 (最后10行):"
                sudo docker logs --tail 10 $container_name 2>&1 || echo "无法获取日志"

                echo -e "\n=========================================="
                echo "输入 'r' 刷新，输入 'q' 退出:"
                ;;
            q|Q)
                echo "退出监控"
                break
                ;;
            *)
                echo "无效输入，请输入 'r' 刷新或 'q' 退出:"
                ;;
        esac
    done
}

# 监控Docker构建进程
monitor_build() {
    echo "监控Docker构建进程 (输入 'r' 刷新，输入 'q' 退出)"
    echo "----------------------------------------"

    # 首次显示信息
    clear
    # 显示当前时间
    echo "监控时间: $(date)"
    echo "----------------------------------------"

    # 显示Docker进程
    echo "Docker相关进程:"
    ps aux | grep -E 'docker (build|compose)' | grep -v grep || echo "当前没有Docker构建进程"

    # 显示系统资源
    echo -e "\n系统资源使用情况:"
    if command -v top &> /dev/null; then
        # macOS和Linux都支持top命令，但参数略有不同
        if [[ "$OSTYPE" == "darwin"* ]]; then
            # macOS
            top -l 1 -n 10 | head -n 15
        else
            # Linux
            top -b -n 1 | head -n 15
        fi
    else
        echo "top命令不可用"
    fi

    # 显示内存使用情况
    echo -e "\n内存使用情况:"
    if command -v free &> /dev/null; then
        free -h
    elif command -v vm_stat &> /dev/null; then
        # macOS
        vm_stat | head -n 10
    else
        echo "无法获取内存使用情况"
    fi

    # 显示Docker磁盘使用情况
    echo -e "\nDocker磁盘使用:"
    if sudo docker info >/dev/null 2>&1; then
        sudo docker system df 2>/dev/null || echo "无法获取Docker磁盘使用情况"
    else
        echo "❌ Docker服务未启动，请先启动Docker Desktop"
    fi

    echo -e "\n=========================================="
    echo "输入 'r' 刷新，输入 'q' 退出:"

    while true; do
        read -n 1 user_input
        echo  # 换行

        case $user_input in
            r|R)
                clear
                # 显示当前时间
                echo "监控时间: $(date)"
                echo "----------------------------------------"

                # 显示Docker进程
                echo "Docker相关进程:"
                ps aux | grep -E 'docker (build|compose)' | grep -v grep || echo "当前没有Docker构建进程"

                # 显示系统资源
                echo -e "\n系统资源使用情况:"
                if command -v top &> /dev/null; then
                    # macOS和Linux都支持top命令，但参数略有不同
                    if [[ "$OSTYPE" == "darwin"* ]]; then
                        # macOS
                        top -l 1 -n 10 | head -n 15
                    else
                        # Linux
                        top -b -n 1 | head -n 15
                    fi
                else
                    echo "top命令不可用"
                fi

                # 显示内存使用情况
                echo -e "\n内存使用情况:"
                if command -v free &> /dev/null; then
                    free -h
                elif command -v vm_stat &> /dev/null; then
                    # macOS
                    vm_stat | head -n 10
                else
                    echo "无法获取内存使用情况"
                fi

                # 显示Docker磁盘使用情况
                echo -e "\nDocker磁盘使用:"
                if sudo docker info >/dev/null 2>&1; then
                    sudo docker system df 2>/dev/null || echo "无法获取Docker磁盘使用情况"
                else
                    echo "❌ Docker服务未启动，请先启动Docker Desktop"
                fi

                echo -e "\n=========================================="
                echo "输入 'r' 刷新，输入 'q' 退出:"
                ;;
            q|Q)
                echo "退出监控"
                break
                ;;
            *)
                echo "无效输入，请输入 'r' 刷新或 'q' 退出:"
                ;;
        esac
    done
}

# ===========================================================
# 主程序
# ===========================================================
main() {
    if [ "$1" == "local" ]; then
        run_local "$@"
        return
    elif [ "$1" == "docker" ]; then
        run_docker
        return
    fi

    while true; do
        show_menu
        read choice
        case $choice in
            1) run_local "$@"; break ;;
            2) run_docker; break ;;
            3) build_only; break ;;
            4) start_service; break ;;
            5) stop_docker ;;
            6) show_logs ;;
            7) rebuild_all; break ;;
            8) fix_docker; break ;;
            9) install_code_check_tools; break ;;
            10) run_code_check; break ;;
            11) check_python_environment; break ;;
            12) check_gpu; break ;;
            13) monitor_container; break ;;
            14) monitor_build; break ;;
            0) echo "退出脚本"; exit 0 ;;
            *) echo "无效选择，请重试" ;;
        esac
    done
}

# 执行主程序
main "$@"
